"""
PlumeSoft 量化代码库
作者：PlumeSoft
用途：贝叶斯择时炼丹炉
"""

import json
import numbers
import os
import random
import shutil
import pandas as pd
import numpy as np
from optuna import Trial
from datetime import datetime
from program.plumesoft_bys.bys_common import dbg, log, init_env
from program import Config as Cfg
from program import Functions as Fun

import warnings

warnings.filterwarnings("ignore")

################################
# 炉子公共处理函数
################################


def replace_macro_path(params: dict, path: str):
    now = datetime.now()
    replace_dict = {
        "$STGNAME$": params["base_stg_file"],
        "$DATE$": now.strftime("%Y%m%d"),
        "$TIME$": now.strftime("%H%M%S"),
        "$BATCH-ID$": params.get("批次", ""),
        "$ROUND-ID$": params.get("轮次", ""),
        "$ROLLING-ID$": params.get("滚动号", ""),
    }
    for key, val in replace_dict.items():
        path = path.replace(key, val)
    return path


# 初始化用于炼丹的策略和全局参数
def init_bys_stg_and_params(params: dict):
    # 更新数据文件路径
    if "批次" in params.keys():
        params["work_path"] = params["batch_work_path"]
        params["finished_path"] = params["batch_finished_path"]
    else:
        params["work_path"] = params["single_work_path"]
        params["finished_path"] = params["single_finished_path"]
    work_path = os.path.join(
        params["data_path"], replace_macro_path(params, params["work_path"])
    )
    os.makedirs(work_path, exist_ok=True)
    old_log_file = params.get("study_log_file", None)
    for key in [
        "study_params_swap_pkl",
        "study_storage_file",
        "study_result_pkl",
        "study_result_csv",
        "study_result_txt",
        "study_factor_params_csv",
        "study_batch_cache_pkl",
        "study_log_file",
    ]:
        if params.get(key, None) is not None:
            params[key] = os.path.join(work_path, os.path.basename(params[key]))

    # 把初始化目录之前的日志文件移动过来
    new_log_file = params.get("study_log_file", None)
    if old_log_file != new_log_file and os.path.isfile(old_log_file):
        if not os.path.isfile(new_log_file):
            shutil.move(old_log_file, new_log_file)
        else:
            # 打开旧日志，读取其内容
            with open(old_log_file, "r", encoding="utf-8") as source:
                source_content = source.read()

            # 打开目标文件，以追加模式
            with open(new_log_file, "a", encoding="utf-8") as target:
                # 将源文件内容追加到目标文件末尾
                target.write(source_content)
            os.remove(old_log_file)

    init_env(params)

    # =====动态导入选币策略
    cls = __import__(
        f'program.{params["base_stg_folder"]}.{params["base_stg_file"]}',
        fromlist=("",),
    )
    log("策略文件：", f'{params["base_stg_folder"]}.{params["base_stg_file"]}')

    stg_params = {}
    for key in params:
        if hasattr(cls, key) and getattr(cls, key) is not None:
            params[key] = getattr(cls, key)
            stg_params[key] = params[key]
    log(f"从策略文件中更新配置参数: {str(stg_params)}")
    dbg(f"更新策略配置参数: {params}")

    # 处理分段日期列表
    if "sample_date_list" in params.keys():
        if params["sample_date_list"] is None or len(params["sample_date_list"]) == 0:
            params["sample_date_list"] = [
                (params["sample_date_start"], params["sample_date_end"])
            ]
            log(
                f'样本区间：{params["sample_date_start"]} - {params["sample_date_end"]}'
            )
        else:
            date_start = (
                min([pd.to_datetime(x[0]) for x in params["sample_date_list"]])
                .date()
                .strftime("%Y-%m-%d")
            )
            date_end = (
                max(
                    [
                        pd.to_datetime(x[1] if x[1] is not None else datetime.today())
                        for x in params["sample_date_list"]
                    ]
                )
                .date()
                .strftime("%Y-%m-%d")
            )
            params["sample_date_start"] = date_start
            params["sample_date_end"] = date_end
            log(f'样本区间：{params["sample_date_list"]}')

    log(f'并行任务：{params["study_n_jobs"]}')

    return cls, params


# 获取因子列表对应的因子列名列表
def get_factor_params_and_column_list(factor_list, param_list, filename="factors"):
    factor_param_dict = {}
    column_list = []
    # 根据config中设置的因子列表，逐个计算每个因子的数据
    for factor in factor_list:
        factor = factor.split(".")[0]
        _cls = __import__("%s.%s" % (filename, factor), fromlist=("",))
        # 兼容get_parameter
        if "get_parameter" in _cls.__dict__:
            # 如果存在get_parameter，直接覆盖config中配置的factor_param_list
            _param_list = getattr(_cls, "get_parameter")()
        else:  # 如果不存在get_parameter，直接用config中配置的factor_param_list
            _param_list = param_list.copy()
        factor_param_dict[factor] = _param_list
        # 遍历参数，计算每个参数对应的因子值
        for n in _param_list:
            # print(f"计算因子:{factor}_{str(n)}")
            # 构建因子列名
            factor_name = f"{factor}_{str(n)}"
            column_list.append(factor_name)
    return factor_param_dict, sorted(list(set(column_list)))


def dict_to_json_str(dt):
    try:
        dt = str(dt)
        # 导出前处理
        for k, v in {
            "nan": "None",
            "NaN": "None",
            "NaT": "None",
            "null": "None",
            "-inf": "None",
            "inf": "None",
            "false": "False",
            "true": "True",
        }.items():
            dt = dt.replace(k, v)

        # 导出为 json 格式
        res = json.dumps(eval(dt), indent=4, ensure_ascii=False)

        # 导出后处理
        for key, value in {
            ": true": ": True",
            ": false": ": False",
            "null": "None",
        }.items():
            res = res.replace(key, value)
        return res
    except:
        return str(dt)


def list_to_code(items):
    res = "[\n" + "\n".join([f"    {str(item)}," for item in items]) + "\n]"
    return res


def move_files_to_finished_path(params: dict, stg_file_list: list):
    # 删除临时文件
    if (
        params.get("batch_last", True)
        and params["study_params_swap_pkl"] is not None
        and os.path.isfile(params["study_params_swap_pkl"])
    ):
        os.remove(params["study_params_swap_pkl"])

    # 批量最后一轮删除缓存文件
    if (
        params.get("batch_last", False)
        and (params["study_batch_cache_pkl"] is not None)
        and os.path.isfile(params["study_batch_cache_pkl"])
    ):
        os.remove(params["study_batch_cache_pkl"])

    # 数据文件移到子目录中，这样可以继续下一轮新的炼丹
    finished_path = os.path.join(
        params["data_path"], replace_macro_path(params, params["finished_path"])
    )
    os.makedirs(finished_path, exist_ok=True)
    for name in [
        "study_storage_file",
        "study_result_pkl",
        "study_result_csv",
        "study_result_txt",
        "study_factor_params_csv",
        "study_log_file",
    ]:
        if os.path.isfile(params[name]):
            shutil.move(
                params[name],
                os.path.join(
                    finished_path,
                    os.path.basename(params[name]),
                ),
            )

    # 复制一份输出的策略文件存档
    out_stg_path = os.path.join(finished_path, "结果策略")
    for stg_file in stg_file_list:
        os.makedirs(out_stg_path, exist_ok=True)
        fn = os.path.join(
            Cfg.root_path, "program", params["output_stg_folder"], f"{stg_file}.py"
        )
        if os.path.isfile(fn):
            shutil.copyfile(fn, os.path.join(out_stg_path, os.path.basename(fn)))


################################
# 过滤因子相关函数
################################


# 取得过滤查询条件列表
def suggest_filter_list(
    trial: Trial, params: dict, mode, factor_param_dict, prefix=None
):
    # 生成过滤查询条件列表
    bys_factor_min_count = params[f"bys_{mode}_factor_min_count"]
    bys_factor_max_count = params[f"bys_{mode}_factor_max_count"]
    if bys_factor_min_count is None or bys_factor_min_count < 1:
        bys_factor_min_count = 1
    if prefix is None:
        prefix = mode

    filter_factor_list = []

    # 为了限制使用因子的数据，先建一个随机因子列表，然后根据这个列表，在限制数量范围内选择使用哪些因子
    # 贝叶斯的参数名称不支持中文，使用固定的序号 n 代替
    tmp_list = [[n, item] for n, item in enumerate(factor_param_dict.items())]
    random.shuffle(tmp_list)

    # 遍历所有因子
    for idx, (n, (factor, params)) in enumerate(tmp_list):
        if bys_factor_min_count is not None and len(
            tmp_list
        ) - idx <= bys_factor_min_count - len(filter_factor_list):
            # 剩余因子不足时直接使用
            use_factor = True
        elif (
            bys_factor_max_count is not None
            and len(filter_factor_list) >= bys_factor_max_count
        ):
            # 因子数量已满时不再使用
            use_factor = False
        else:
            # 最后由贝叶斯框架决定是否启用
            use_factor = trial.suggest_categorical(
                f"{prefix}.{factor}.启用", [False, True]
            )

        if use_factor:
            # 取得因子及参数
            if len(params) == 1:
                param = params[0]
            elif len(params) > 1:
                param = trial.suggest_categorical(f"{prefix}.{factor}.参数", params)
            else:
                param = 0
            filter_factor_list.append((factor, param))

    return filter_factor_list


################################
# 排序因子相关函数
################################


# 取得排序参数字典
def suggest_sort_factor_list(
    trial: Trial, params: dict, factor_param_dict, factor_exclude_list
):
    # 生成排序参数字典
    sort_factor_list = []

    bys_sort_factor_min_count = params["bys_sort_factor_min_count"]
    bys_sort_factor_max_count = params["bys_sort_factor_max_count"]
    bys_sort_min_power = params["bys_sort_min_power"]
    bys_sort_max_power = params["bys_sort_max_power"]
    if bys_sort_factor_min_count is None or bys_sort_factor_min_count < 1:
        bys_sort_factor_min_count = 1

    # 为了限制使用因子的数据，先建一个随机因子列表，然后根据这个列表，在限制数量范围内选择使用哪些因子
    # 贝叶斯的参数名称不支持中文，使用固定的序号 n 代替
    tmp_list = [[n, item] for n, item in enumerate(factor_param_dict.items())]
    random.shuffle(tmp_list)

    # 遍历所有因子
    for idx, (n, (factor, params)) in enumerate(tmp_list):
        if bys_sort_factor_min_count is not None and len(
            tmp_list
        ) - idx <= bys_sort_factor_min_count - len(sort_factor_list):
            # 剩余因子不足时直接使用
            use_factor = True
        elif factor in factor_exclude_list:
            # 因子不能在排除列表中
            use_factor = False
        elif (
            bys_sort_factor_max_count is not None
            and len(sort_factor_list) >= bys_sort_factor_max_count
        ):
            # 因子数量已满时不再使用
            use_factor = False
        else:
            # 最后由贝叶斯框架决定是否启用
            use_factor = trial.suggest_categorical(f"排序.{factor}.启用", [False, True])

        if use_factor:
            # 取得因子及参数
            if len(params) == 1:
                param = params[0]
            elif len(params) > 1:
                param = trial.suggest_categorical(f"排序.{factor}.参数", params)
            else:
                param = 0
            sort_asc = trial.suggest_categorical(
                f"排序.{factor}.排序方式", [False, True]
            )
            if bys_sort_min_power <= bys_sort_max_power + 1:
                sort_power = trial.suggest_int(
                    f"排序.{factor}.排序权重",
                    bys_sort_min_power,
                    bys_sort_max_power,
                    step=1,
                )
            else:
                sort_power = bys_sort_min_power
            sort_factor_list.append((factor, sort_asc, param, sort_power))

    return sort_factor_list


################################
# 用户参数相关函数
################################


# 取得用户参数字典
def do_suggest_user_params(
    trial: Trial, params_dict: dict, prefix="用户", macro_dict={}
):
    ret_dt = {}
    for pm_key, val in params_dict.items():
        key_list = []
        # 处理参数名中的宏，为每个因子生成一个参数
        for macro, val_list in macro_dict.items():
            if macro in pm_key:
                # 只支持一个宏，否则组合可能性就太多了
                key_list += [pm_key.replace(macro, v) for v in val_list]
                break
        if len(key_list) == 0:
            # 没有宏时直接使用
            key_list = [pm_key]

        # 生成所有参数
        for key in key_list:
            if "枚举列表" in val.keys():
                enum_list = val.get("枚举列表", [])
                ret_dt[key] = trial.suggest_categorical(
                    f"{prefix}.{key}.枚举", enum_list
                )
            else:
                low, high, step = val["下限"], val["上限"], val["步长"]
                if (
                    isinstance(low, int)
                    and isinstance(high, int)
                    and isinstance(step, int)
                ):
                    ret_dt[key] = trial.suggest_int(
                        f"{prefix}.{key}", low, high, step=step
                    )
                else:
                    ret_dt[key] = trial.suggest_float(
                        f"{prefix}.{key}", low, high, step=step
                    )
    return ret_dt


# 取得用户参数字典
def suggest_user_params(
    cls,
    trial: Trial,
    params_dict: dict,
    filter_factors,
    sort_factors,
    prefix="用户",
    macro_dict={},
):
    ret_dt = do_suggest_user_params(trial, params_dict, prefix)
    if hasattr(cls, "cal_extra_user_param_dict"):
        for level in range(1, 4, 1):
            extra_params_dt = cls.cal_extra_user_param_dict(
                ret_dt, filter_factors, sort_factors, level
            )
            if extra_params_dt is None:
                break
            extra_dt = do_suggest_user_params(
                trial, extra_params_dt, prefix, macro_dict
            )
            ret_dt.update(extra_dt)
    return ret_dt


################################
# 评分相关函数
################################


def is_number(n):
    # Python 中 True 和 False 也会识别为 int 类型，通过转为字符串来判断
    return (not str(n) in ["False", "True"]) and (
        isinstance(n, int) or isinstance(n, float)
    )


def cal_segmented_score(points: list, x_target):
    # 分段线性插值算法计算分数
    # 确保points按x值升序排列
    points.sort(key=lambda p: p[0])

    # 查找x_target所在的区间
    for i in range(len(points) - 1):
        x1, y1 = points[i]
        x2, y2 = points[i + 1]
        if (
            (x1 <= x_target <= x2)
            or (i == 0 and x_target <= x1)
            or (i == len(points) - 2 and x_target >= x2)
        ):
            # 使用两点式直线方程 y = y1 + ((y2-y1)/(x2-x1))*(x-x1)
            slope = (y2 - y1) / (x2 - x1)
            y_target = y1 + slope * (x_target - x1)
            return y_target

    return 0


def cal_rtn_score(
    params: dict,
    cls,
    tp,
    eval_dt: dict,
    equity: pd.DataFrame,
    rtn,
    year_return,
    month_return,
    is_sample,
):
    # 计算存档字段
    r_dt = {}
    for col in params["save_rtn_cols"]:
        r_dt[col] = rtn[col].iloc[0]
    year_df = year_return.reset_index()
    year_df["年份"] = year_df["candle_begin_time"].dt.year
    for rec in year_df.to_records():
        r_dt[rec["年份"]] = rec["涨跌幅"]

    # 计算综合分数
    c_dt = {}
    score_select = 0
    invalid_score = False
    for key, dt in eval_dt.items():
        val = float(str(rtn[key].iloc[0]).replace("%", ""))
        c_dt[key] = val
        h_val = dt["高值"]
        h_score = dt["高分"]
        l_val = dt["低值"]
        l_score = dt["低分"]
        extra_points = dt.get("附加", None)
        bys_eval_score_check_low_limit = dt.get(
            "低于低值惩罚", params["bys_eval_score_check_low_limit"]
        )
        bys_eval_score_check_high_limit = dt.get(
            "高于高值惩罚", params["bys_eval_score_check_high_limit"]
        )
        # 有些因子没有值时，不计算评分
        if isinstance(val, numbers.Number) and not np.isinf(val):
            if extra_points is not None:
                score = cal_segmented_score(
                    [(l_val, l_score), (h_val, h_score)] + extra_points, val
                )
            else:
                score = (val - l_val) / (h_val - l_val) * (h_score - l_score) + l_score

            # 检查评分是否低于最低值
            if is_number(bys_eval_score_check_low_limit):
                if score < l_score:
                    score += bys_eval_score_check_low_limit
            elif bys_eval_score_check_low_limit == True:
                if score < l_score:
                    invalid_score = True
            # 检查评分是否高于最高值
            if is_number(bys_eval_score_check_high_limit):
                if score > h_score:
                    score = bys_eval_score_check_high_limit
            elif bys_eval_score_check_high_limit == True:
                if score > h_score:
                    invalid_score = True
        else:
            score = 0
        c_dt[f"{key}_得分"] = round(score, 2)
        score_select += score

    # 只要有一项评分低于最低值，则返回无效评分
    if invalid_score:
        score_select = -2000

    # 调用策略自定义的评分
    if hasattr(cls, "user_cal_score2"):
        score_select, r_dt, c_dt = getattr(cls, "user_cal_score2")(
            score_select,
            tp,
            equity,
            rtn,
            year_return,
            month_return,
            r_dt,
            c_dt,
            is_sample,
        )
    elif hasattr(cls, "user_cal_score"):
        score_select, r_dt, c_dt = getattr(cls, "user_cal_score")(
            score_select,
            tp,
            rtn,
            year_return,
            month_return,
            r_dt,
            c_dt,
            is_sample,
        )
    return round(score_select, 2), r_dt, c_dt
